import os
import paddle

models_path = r'output/cascade_rcnn_res2net_r50_fpn_1x_coco_bfp_blio_nofoal_mratio_best'

models = os.listdir(models_path)

model_item = ['9','10','11','12']
model_pdparams = []
for m in models:
    if m.split('.')[-1] == 'pdparams' and m.split('.')[0] in model_item:
        model_pdparams.append(m)

# print(model_pdparams)
paths = []
for model in model_pdparams:
    paths.append(os.path.join(models_path, model))


models = [paddle.load(p) for p in paths]
model_num = len(models)
param_state_dict = models[-1]
new_state_dict = param_state_dict.copy()


for key in param_state_dict:
    print('key :', key)
    sum_weight = 0.0
    for m in models:
        sum_weight += m[key]
    avg_weight = sum_weight / model_num
    new_state_dict[key] = avg_weight

print('save swa model...')
save_path = os.path.join(models_path, "model_swa.pdparams")
paddle.save(new_state_dict, save_path)
